# -*- encoding: utf-8 -*-
'''
@File    :   jiangnan_reader.py
@Time    :   2022/3/23 8:59
@Author  :   ZhangChaoYang
@Desc    :   把江南大学原始数据文件，经过各种转换（FFT、CWT等）变成numpy ndarray，存入文件。
'''

import os, sys

sys.path.insert(0, os.getcwd())
import re
from data_process import mat_reader
from data_process.mat_reader import convertDir1d, convertDir2d, take_a_look_1d, take_a_look_2d
from absl import flags
from absl import app


FLAGS = flags.FLAGS

flags.DEFINE_string('source_data', "E:\zcy\数据集\JiangNan", '原始数据存在目录')


regex = re.compile("[a-z]{1,2}(\d+)_.*")
fs = 50e3  # 采样频率50kHz
WC = ["600", "800", "1000"]  # 转速代表了工况work condition


class JNUReader(mat_reader.MatFileReader):
    def __init__(self, infile, trans1d="", trans2d="", work_load=""):
        super(JNUReader, self).__init__(infile, "jiangnan", trans1d, trans2d, work_load)

    def get_label(self):
        basename = self.get_basename()
        if basename.lower().startswith("n") > 0:  # 文件名以n开头是正常样本
            return 1
        else:
            return 0

    def window_size(self):
        return 1024

    def read1d(self):
        basename = self.get_basename()
        if not basename.endswith(".csv"):
            return None
        if self.work_load:
            rpm = regex.findall(basename)
            if len(rpm) == 0 or rpm[0] != self.work_load:
                return None
        array = []
        with open(self.infile, encoding="utf-8") as fin:
            lines = fin.readlines()
            for line in lines:
                content = line.strip()
                if content:
                    array.append(float(content))
        x = self.trans1d(array, 1, fs)
        return x

    def read2d(self):
        basename = self.get_basename()
        if not basename.endswith(".csv"):
            return None
        if self.work_load:
            rpm = regex.findall(basename)
            if len(rpm) == 0 or rpm[0] != self.work_load:
                return None
        array = []
        with open(self.infile, encoding="utf-8") as fin:
            lines = fin.readlines()
            for line in lines:
                content = line.strip()
                if content:
                    array.append(float(content))
        x = self.trans2d(array, 1, fs)
        return x

def main(argv):
    for trans1d in ['original', 'stat', 'fft']:
        reader = JNUReader("", trans1d=trans1d)
        convertDir1d(reader, [FLAGS.source_data])
        for wc in WC:
            reader = JNUReader("", trans1d=trans1d, work_load=wc)
            convertDir1d(reader, [FLAGS.source_data])
        take_a_look_1d(reader.name, trans1d)

    for trans2d in ['sfft', 'cwt']:
        reader = JNUReader(FLAGS.source_data, trans2d=trans2d)
        convertDir2d(reader, [FLAGS.source_data])
        for wc in WC:
            reader = JNUReader("", trans2d=trans2d, work_load=wc)
            convertDir2d(reader, [FLAGS.source_data])
        take_a_look_2d(reader.name, trans2d)


if __name__ == '__main__':
    app.run(main)

# python .\data_process\jiangnan_reader.py --source_data E:\zcy\数据集\JiangNan